package com.aikg.kgrag_java.controller;

import com.aikg.kgrag_java.common.Result;
import com.aikg.kgrag_java.pojo.entity.ConversationInfo;
import com.aikg.kgrag_java.pojo.vo.ChunkInfoVO;
import com.aikg.kgrag_java.pojo.vo.Neo4jRetrievalVO;
import com.aikg.kgrag_java.service.IStreamService;
import com.aikg.kgrag_java.service.IntentRecognitionService;
import com.aikg.kgrag_java.service.impl.QuestionRecommendPipeline;
import com.aikg.kgrag_java.service.impl.RAGPipelineService;
import com.github.benmanes.caffeine.cache.Cache;
import jakarta.annotation.Resource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.UUID;

/**
 * @ClassName: ChatController
 * @Description: 对话相关的接口类
 * @Author: 李宗雨
 * @Date: 2025/1/6 10:11
 * @Version: 1.0
 **/

@RestController
public class ChatController {

    @Autowired
    QuestionRecommendPipeline questionRecommendPipeline;

    @Autowired
    IStreamService streamService;

    @Autowired
    RAGPipelineService ragPipeline;

    @Resource(name = "conversationCache")
    Cache<String, ConversationInfo> conversationCache;

    @Autowired
    IntentRecognitionService intentRecognitionService;

    @GetMapping("/stream_chat")
    public SseEmitter streamChat(int projectId, String query,int top_k, float score) throws IOException {
        String conversationId = UUID.randomUUID().toString();
        System.out.println(conversationId);
//        // 意图识别
//        List<String> intent = intentRecognitionService.IntentRecognition(query);
//        System.out.println(conversationId);
//        return switch (intent.get(0)) {
//            // 模板
//            case "1" -> streamService.streamByContent(intent.get(1), conversationId);
//            // RAG
//            case "2" -> ragPipeline.rag(query, top_k, projectId, conversationId);
//            // 查水文
//            // 查气象
//            case "3", "4" -> streamService.streamByContent(intent.get(2), conversationId);
//            // 不可能走到这里，放心，为了防止报错
//            default -> new SseEmitter();
//        };
        return ragPipeline.rag(query, top_k, projectId, conversationId);
    }

    /**
     * 获取某对话的文本检索结果
     * @param conversationId: 对话的id
     * @return 相关文本块列表
     */
    @GetMapping("/get_chunk")
    public Result<List<ChunkInfoVO>> getChunk(@RequestParam(name = "conversation_id") String conversationId){
        if (conversationCache.getIfPresent(conversationId)!=null){
            return Result.buildSuccessResultWithData(conversationCache.getIfPresent(conversationId).getChunks());
        }
        return Result.buildSuccessResult("文本为空", new ArrayList<>());
    }

    /**
     * 获取某对话的图数据库检索结果
     * @param conversationId: 对话的id
     * @return 相关子图数据
     */
    @GetMapping("/get_neo4j_data")
    public Result<Neo4jRetrievalVO> getNeo4jData(@RequestParam(name = "conversation_id") String conversationId){
        if (conversationCache.getIfPresent(conversationId)!=null) {
            return Result.buildSuccessResult("Success", conversationCache.getIfPresent(conversationId).getNeo4jRetrievalVO());
        }
        return Result.buildSuccessResultWithData(null);
    }

    /**
     * 问题推荐
     *
     * @param conversationId 问答对的id
     * @return 问题数组
     */
    @PostMapping("/questionRecommend")
    public Result<List<String>> questionRecommendApi(String conversationId){
        ConversationInfo conversationInfo = conversationCache.getIfPresent(conversationId);
        if (conversationInfo != null){
            return Result.buildSuccessResultWithData(
                    questionRecommendPipeline.questionRecommend(conversationInfo, 3)
            );
        }
        return Result.buildSuccessResult("问题为空", new ArrayList<String>());
    }
}
